"""
MCP相关的Pydantic schemas

用于API请求和响应的数据验证
"""

from pydantic import BaseModel, Field, ConfigDict
from typing import Optional, Dict, Any, List
from datetime import datetime
from enum import Enum


# ==================== 枚举类型 ====================

class SessionStatus(str, Enum):
    """会话状态"""
    CONNECTING = "connecting"
    CONNECTED = "connected"
    DISCONNECTED = "disconnected"
    ERROR = "error"


class ContextType(str, Enum):
    """上下文类型"""
    SYSTEM = "system"
    USER = "user"
    ASSISTANT = "assistant"
    TOOL = "tool"


class MessageType(str, Enum):
    """消息类型"""
    CONNECT = "CONNECT"
    CONNECT_ACK = "CONNECT_ACK"
    REQUEST = "REQUEST"
    RESPONSE = "RESPONSE"
    ERROR = "ERROR"
    PING = "PING"
    PONG = "PONG"
    DISCONNECT = "DISCONNECT"


# ==================== MCP会话相关 ====================

class MCPSessionCreate(BaseModel):
    """创建MCP会话的请求"""
    mcp_server_url: Optional[str] = Field(None, description="MCP服务器地址（可选，使用默认配置）")
    conversation_id: Optional[int] = Field(None, description="关联的对话ID")
    model_config_id: Optional[int] = Field(None, description="关联的模型配置ID")
    initial_context: Optional[Dict[str, Any]] = Field(None, description="初始上下文数据")
    connection_config: Optional[Dict[str, Any]] = Field(None, description="连接配置")


class MCPSessionResponse(BaseModel):
    """MCP会话响应"""
    id: int = Field(..., description="会话数据库ID")
    session_id: str = Field(..., description="MCP会话ID（UUID）")
    user_id: int = Field(..., description="用户ID")
    mcp_server_url: str = Field(..., description="MCP服务器地址")
    conversation_id: Optional[int] = Field(None, description="关联的对话ID")
    model_config_id: Optional[int] = Field(None, description="关联的模型配置ID")
    status: SessionStatus = Field(..., description="会话状态")
    is_active: bool = Field(..., description="是否活跃")
    message_count: int = Field(..., description="消息总数")
    error_count: int = Field(..., description="错误次数")
    last_heartbeat_at: Optional[datetime] = Field(None, description="最后心跳时间")
    connected_at: Optional[datetime] = Field(None, description="连接建立时间")
    disconnected_at: Optional[datetime] = Field(None, description="断开连接时间")
    created_at: datetime = Field(..., description="创建时间")
    updated_at: datetime = Field(..., description="更新时间")
    
    model_config = ConfigDict(from_attributes=True)


class MCPSessionUpdate(BaseModel):
    """更新MCP会话"""
    status: Optional[SessionStatus] = Field(None, description="会话状态")
    is_active: Optional[bool] = Field(None, description="是否活跃")
    connection_config: Optional[Dict[str, Any]] = Field(None, description="连接配置")


# ==================== MCP上下文相关 ====================

class MCPContextCreate(BaseModel):
    """创建MCP上下文的请求"""
    session_id: str = Field(..., description="MCP会话ID")
    context_data: Dict[str, Any] = Field(..., description="上下文数据")
    context_type: ContextType = Field(ContextType.USER, description="上下文类型")
    model_id: Optional[str] = Field(None, description="模型ID")
    context_metadata: Optional[Dict[str, Any]] = Field(None, description="元数据")


class MCPContextResponse(BaseModel):
    """MCP上下文响应"""
    id: int = Field(..., description="上下文ID")
    session_id: str = Field(..., description="MCP会话ID")
    model_id: Optional[str] = Field(None, description="模型ID")
    context_data: Dict[str, Any] = Field(..., description="上下文数据")
    context_type: ContextType = Field(..., description="上下文类型")
    context_metadata: Optional[Dict[str, Any]] = Field(None, description="元数据")
    version: int = Field(..., description="版本号")
    token_count: int = Field(..., description="Token数量")
    size_bytes: int = Field(..., description="数据大小（字节）")
    is_encrypted: bool = Field(..., description="是否加密")
    created_at: datetime = Field(..., description="创建时间")
    updated_at: datetime = Field(..., description="更新时间")
    
    model_config = ConfigDict(from_attributes=True)


class MCPContextUpdate(BaseModel):
    """更新MCP上下文"""
    context_data: Dict[str, Any] = Field(..., description="上下文数据")
    increment_version: bool = Field(True, description="是否增加版本号")


# ==================== MCP消息相关 ====================

class MCPMessageCreate(BaseModel):
    """创建MCP消息的请求"""
    session_id: str = Field(..., description="MCP会话ID")
    message_id: str = Field(..., description="消息ID（UUID）")
    message_type: MessageType = Field(..., description="消息类型")
    direction: str = Field(..., description="消息方向（send/receive）")
    payload: Dict[str, Any] = Field(..., description="消息负载")


class MCPMessageResponse(BaseModel):
    """MCP消息响应"""
    id: int = Field(..., description="消息数据库ID")
    session_id: str = Field(..., description="MCP会话ID")
    message_id: str = Field(..., description="消息ID（UUID）")
    message_type: MessageType = Field(..., description="消息类型")
    direction: str = Field(..., description="消息方向")
    payload: Dict[str, Any] = Field(..., description="消息负载")
    status: str = Field(..., description="消息状态")
    error_message: Optional[str] = Field(None, description="错误信息")
    send_at: Optional[datetime] = Field(None, description="发送时间")
    receive_at: Optional[datetime] = Field(None, description="接收时间")
    duration_ms: Optional[int] = Field(None, description="处理时长（毫秒）")
    created_at: datetime = Field(..., description="创建时间")
    
    model_config = ConfigDict(from_attributes=True)


# ==================== MCP请求/响应 ====================

class MCPRequestPayload(BaseModel):
    """MCP请求负载"""
    payload: Dict[str, Any] = Field(..., description="请求负载数据")
    update_context: bool = Field(True, description="是否更新上下文")
    timeout: Optional[int] = Field(None, description="超时时间（秒）")


class MCPResponsePayload(BaseModel):
    """MCP响应负载"""
    response: Dict[str, Any] = Field(..., description="响应数据")
    context_updated: bool = Field(..., description="上下文是否已更新")
    message_id: str = Field(..., description="消息ID")
    duration_ms: Optional[int] = Field(None, description="处理时长（毫秒）")


# ==================== 统计信息 ====================

class MCPSessionStats(BaseModel):
    """MCP会话统计信息"""
    total_sessions: int = Field(..., description="总会话数")
    active_sessions: int = Field(..., description="活跃会话数")
    total_messages: int = Field(..., description="总消息数")
    total_contexts: int = Field(..., description="总上下文数")
    avg_message_duration_ms: Optional[float] = Field(None, description="平均消息处理时长（毫秒）")
    error_rate: float = Field(..., description="错误率")

